{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 7: Large-Scale Simulations\n",
    "\n",
    "**Duration:** ~35 minutes | **Prerequisites:** All Basic Tutorials\n",
    "\n",
    "## Learning Objectives\n",
    "\n",
    "By the end of this tutorial, you will:\n",
    "\n",
    "- ✅ Optimize memory usage for large networks\n",
    "- ✅ Apply JIT compilation best practices\n",
    "- ✅ Use batching strategies effectively\n",
    "- ✅ Leverage GPU/TPU acceleration\n",
    "- ✅ Profile and optimize performance\n",
    "- ✅ Implement sparse connectivity\n",
    "\n",
    "## Overview\n",
    "\n",
    "Scaling neural simulations to thousands or millions of neurons requires careful optimization. BrainPy leverages JAX for high-performance computing on CPUs, GPUs, and TPUs.\n",
    "\n",
    "**Key concepts:**\n",
    "- **JIT compilation**: Compile Python code to optimized machine code\n",
    "- **Memory efficiency**: Minimize state storage and intermediate computations\n",
    "- **Sparse operations**: Only compute where connections exist\n",
    "- **Batching**: Process multiple trials simultaneously\n",
    "- **Device acceleration**: Utilize GPU/TPU parallelism\n",
    "\n",
    "Let's learn how to build efficient large-scale simulations!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import braintools\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "# Set random seed\n",
    "brainstate.random.seed(42)\n",
    "\n",
    "# Configure environment\n",
    "brainstate.environ.set(dt=0.1 * u.ms)\n",
    "\n",
    "# Check available devices\n",
    "print(\"🖥️  Available devices:\")\n",
    "print(f\"   {jax.devices()}\")\n",
    "print(f\"   Default backend: {jax.default_backend()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: JIT Compilation Basics\n",
    "\n",
    "Just-In-Time (JIT) compilation converts Python code to optimized machine code. This can provide 10-100× speedups!\n",
    "\n",
    "**Benefits of JIT:**\n",
    "- Eliminates Python interpreter overhead\n",
    "- Enables compiler optimizations (loop fusion, vectorization)\n",
    "- Required for GPU/TPU execution\n",
    "\n",
    "**Rules for JIT:**\n",
    "- Functions must be pure (no side effects)\n",
    "- Array shapes must be static (known at compile time)\n",
    "- Avoid Python loops over dynamic ranges\n",
    "\n",
    "Let's compare JIT vs non-JIT performance!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple network without JIT\n",
    "class SimpleNetwork(brainstate.nn.Module):\n",
    "    def __init__(self, n_neurons=1000):\n",
    "        super().__init__()\n",
    "        self.neurons = brainpy.state.LIF(\n",
    "            n_neurons,\n",
    "            V_rest=-65.0*u.mV, V_th=-50.0*u.mV, tau=10.0*u.ms\n",
    "        )\n",
    "    \n",
    "    def update(self, inp):\n",
    "        self.neurons(inp)\n",
    "        return self.neurons.get_spike()\n",
    "\n",
    "# Test without JIT\n",
    "net_no_jit = SimpleNetwork(n_neurons=1000)\n",
    "brainstate.nn.init_all_states(net_no_jit)\n",
    "\n",
    "# Warmup\n",
    "inp = brainstate.random.rand(1000) * 2.0 * u.nA\n",
    "_ = net_no_jit(inp)\n",
    "\n",
    "# Time execution\n",
    "n_steps = 1000\n",
    "start = time.time()\n",
    "for _ in range(n_steps):\n",
    "    inp = brainstate.random.rand(1000) * 2.0 * u.nA\n",
    "    _ = net_no_jit(inp)\n",
    "time_no_jit = time.time() - start\n",
    "\n",
    "print(f\"⏱️  Without JIT: {time_no_jit:.3f} seconds for {n_steps} steps\")\n",
    "print(f\"   ({time_no_jit/n_steps*1000:.2f} ms/step)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same network WITH JIT\n",
    "net_jit = SimpleNetwork(n_neurons=1000)\n",
    "brainstate.nn.init_all_states(net_jit)\n",
    "\n",
    "# Apply JIT compilation\n",
    "@brainstate.transform.jit\n",
    "def run_step_jit(net, inp):\n",
    "    return net(inp)\n",
    "\n",
    "# Warmup (compilation happens here)\n",
    "inp = brainstate.random.rand(1000) * 2.0 * u.nA\n",
    "_ = run_step_jit(net_jit, inp)\n",
    "\n",
    "# Time execution\n",
    "start = time.time()\n",
    "for _ in range(n_steps):\n",
    "    inp = brainstate.random.rand(1000) * 2.0 * u.nA\n",
    "    _ = run_step_jit(net_jit, inp)\n",
    "time_jit = time.time() - start\n",
    "\n",
    "print(f\"⏱️  With JIT: {time_jit:.3f} seconds for {n_steps} steps\")\n",
    "print(f\"   ({time_jit/n_steps*1000:.2f} ms/step)\")\n",
    "print(f\"\\n🚀 Speedup: {time_no_jit/time_jit:.1f}×\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Memory Optimization\n",
    "\n",
    "Large networks require careful memory management. Key strategies:\n",
    "\n",
    "1. **Use appropriate data types**: Float32 instead of Float64\n",
    "2. **Minimize state storage**: Only keep necessary variables\n",
    "3. **Avoid unnecessary copies**: Use in-place updates where possible\n",
    "4. **Clear intermediate results**: Don't accumulate large histories\n",
    "\n",
    "Let's compare memory usage for different approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate memory usage\n",
    "def estimate_memory_mb(n_neurons, n_synapses, dtype_bytes=4):\n",
    "    \"\"\"Estimate memory requirements.\n",
    "    \n",
    "    Args:\n",
    "        n_neurons: Number of neurons\n",
    "        n_synapses: Number of synaptic connections\n",
    "        dtype_bytes: Bytes per element (4 for float32, 8 for float64)\n",
    "    \"\"\"\n",
    "    # Neuron states (V, spike)\n",
    "    neuron_memory = n_neurons * 2 * dtype_bytes\n",
    "    \n",
    "    # Synapse states (g, x for plasticity)\n",
    "    synapse_memory = n_synapses * 2 * dtype_bytes\n",
    "    \n",
    "    # Connection weights\n",
    "    weight_memory = n_synapses * dtype_bytes\n",
    "    \n",
    "    total_bytes = neuron_memory + synapse_memory + weight_memory\n",
    "    total_mb = total_bytes / (1024 * 1024)\n",
    "    \n",
    "    return total_mb\n",
    "\n",
    "# Compare different network sizes\n",
    "sizes = [100, 1000, 10000, 100000, 1000000]\n",
    "connectivity = 0.1\n",
    "\n",
    "print(\"📊 Memory Requirements (Float32):\")\n",
    "print(\"=\"*60)\n",
    "print(f\"{'Neurons':<12} {'Synapses':<15} {'Memory (MB)':<15} {'Memory (GB)'}\")\n",
    "print(\"=\"*60)\n",
    "\n",
    "for n in sizes:\n",
    "    n_syn = int(n * n * connectivity)\n",
    "    mem_mb = estimate_memory_mb(n, n_syn, dtype_bytes=4)\n",
    "    mem_gb = mem_mb / 1024\n",
    "    print(f\"{n:<12,} {n_syn:<15,} {mem_mb:<15.2f} {mem_gb:<.3f}\")\n",
    "\n",
    "print(\"\\n💡 Optimization tips:\")\n",
    "print(\"   • Use sparse connectivity to reduce synapse count\")\n",
    "print(\"   • Use float32 instead of float64 (2× memory savings)\")\n",
    "print(\"   • Don't store full spike history (record only what you need)\")\n",
    "print(\"   • Process in batches if memory-constrained\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Sparse Connectivity\n",
    "\n",
    "Biological networks are sparsely connected (~1-10% connectivity). Using sparse matrices dramatically reduces memory and computation.\n",
    "\n",
    "**Dense vs Sparse:**\n",
    "- Dense: Store all $N \\times N$ connections (even zeros)\n",
    "- Sparse: Store only non-zero connections\n",
    "\n",
    "**Memory savings:**\n",
    "- 10% connectivity → 90% memory reduction\n",
    "- 1% connectivity → 99% memory reduction\n",
    "\n",
    "BrainPy's `EventFixedProb` connection automatically uses sparse representations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare dense vs sparse connectivity\n",
    "n_pre = 1000\n",
    "n_post = 1000\n",
    "prob = 0.05  # 5% connectivity\n",
    "\n",
    "# Dense connection matrix\n",
    "dense_matrix = (np.random.rand(n_post, n_pre) < prob).astype(np.float32)\n",
    "dense_size_mb = dense_matrix.nbytes / (1024 * 1024)\n",
    "\n",
    "# Sparse representation (only store indices and values)\n",
    "indices = np.argwhere(dense_matrix > 0)\n",
    "values = dense_matrix[dense_matrix > 0]\n",
    "sparse_size_mb = (indices.nbytes + values.nbytes) / (1024 * 1024)\n",
    "\n",
    "print(\"🔍 Dense vs Sparse Comparison:\")\n",
    "print(f\"   Network size: {n_pre} → {n_post} neurons\")\n",
    "print(f\"   Connectivity: {prob*100}%\")\n",
    "print(f\"   Actual connections: {len(values):,}\")\n",
    "print()\n",
    "print(f\"   Dense storage: {dense_size_mb:.2f} MB\")\n",
    "print(f\"   Sparse storage: {sparse_size_mb:.2f} MB\")\n",
    "print(f\"   Memory savings: {(1 - sparse_size_mb/dense_size_mb)*100:.1f}%\")\n",
    "print(f\"   Space ratio: {dense_size_mb/sparse_size_mb:.1f}×\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build large sparse network\n",
    "class LargeSparseNetwork(brainstate.nn.Module):\n",
    "    \"\"\"Large network with sparse connectivity.\"\"\"\n",
    "    \n",
    "    def __init__(self, n_exc=4000, n_inh=1000, p_conn=0.02):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Neurons\n",
    "        self.E = brainpy.state.LIF(n_exc, V_rest=-65.*u.mV, V_th=-50.*u.mV, tau=15.*u.ms)\n",
    "        self.I = brainpy.state.LIF(n_inh, V_rest=-65.*u.mV, V_th=-50.*u.mV, tau=10.*u.ms)\n",
    "        \n",
    "        # Sparse projections with EventFixedProb\n",
    "        self.E2E = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_exc, n_exc, p_conn, 0.6*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_exc, tau=5.*u.ms),\n",
    "            out=brainpy.state.COBA.desc(E=0.*u.mV),\n",
    "            post=self.E\n",
    "        )\n",
    "        \n",
    "        self.E2I = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_exc, n_inh, p_conn, 0.6*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_inh, tau=5.*u.ms),\n",
    "            out=brainpy.state.COBA.desc(E=0.*u.mV),\n",
    "            post=self.I\n",
    "        )\n",
    "        \n",
    "        self.I2E = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_inh, n_exc, p_conn, 6.7*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_exc, tau=10.*u.ms),\n",
    "            out=brainpy.state.COBA.desc(E=-80.*u.mV),\n",
    "            post=self.E\n",
    "        )\n",
    "        \n",
    "        self.I2I = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_inh, n_inh, p_conn, 6.7*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_inh, tau=10.*u.ms),\n",
    "            out=brainpy.state.COBA.desc(E=-80.*u.mV),\n",
    "            post=self.I\n",
    "        )\n",
    "    \n",
    "    def update(self, inp_e, inp_i):\n",
    "        spk_e = self.E.get_spike()\n",
    "        spk_i = self.I.get_spike()\n",
    "        \n",
    "        self.E2E(spk_e)\n",
    "        self.E2I(spk_e)\n",
    "        self.I2E(spk_i)\n",
    "        self.I2I(spk_i)\n",
    "        \n",
    "        self.E(inp_e)\n",
    "        self.I(inp_i)\n",
    "        \n",
    "        return spk_e, spk_i\n",
    "\n",
    "# Create large network\n",
    "large_net = LargeSparseNetwork(n_exc=4000, n_inh=1000, p_conn=0.02)\n",
    "brainstate.nn.init_all_states(large_net)\n",
    "\n",
    "print(\"✅ Created large sparse network:\")\n",
    "print(f\"   Excitatory neurons: 4,000\")\n",
    "print(f\"   Inhibitory neurons: 1,000\")\n",
    "print(f\"   Total neurons: 5,000\")\n",
    "print(f\"   Connectivity: 2%\")\n",
    "print(f\"   Approximate connections: {5000*5000*0.02:,.0f}\")\n",
    "print(f\"   Estimated memory: ~20 MB (sparse) vs ~400 MB (dense)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Batching for Parallelism\n",
    "\n",
    "Running multiple independent simulations (trials) can be done in parallel using batching. This is especially efficient on GPUs.\n",
    "\n",
    "**Batching benefits:**\n",
    "- Run multiple trials simultaneously\n",
    "- Amortize compilation cost\n",
    "- Better GPU utilization\n",
    "- Faster parameter sweeps\n",
    "\n",
    "**How it works:**\n",
    "- Add batch dimension: `(batch_size, n_neurons)`\n",
    "- Operations automatically vectorized\n",
    "- Each trial independent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Single trial simulation\n",
    "def simulate_single_trial(n_steps=1000):\n",
    "    net = SimpleNetwork(n_neurons=1000)\n",
    "    brainstate.nn.init_all_states(net)\n",
    "    \n",
    "    @brainstate.transform.jit\n",
    "    def step(net, inp):\n",
    "        return net(inp)\n",
    "    \n",
    "    for _ in range(n_steps):\n",
    "        inp = brainstate.random.rand(1000) * 2.0 * u.nA\n",
    "        _ = step(net, inp)\n",
    "\n",
    "# Batched simulation\n",
    "def simulate_batched_trials(n_trials=10, n_steps=1000):\n",
    "    net = SimpleNetwork(n_neurons=1000)\n",
    "    brainstate.nn.init_all_states(net, batch_size=n_trials)\n",
    "    \n",
    "    @brainstate.transform.jit\n",
    "    def step(net, inp):\n",
    "        return net(inp)\n",
    "    \n",
    "    for _ in range(n_steps):\n",
    "        inp = brainstate.random.rand(n_trials, 1000) * 2.0 * u.nA\n",
    "        _ = step(net, inp)\n",
    "\n",
    "# Compare timing\n",
    "n_trials = 10\n",
    "\n",
    "# Sequential trials\n",
    "start = time.time()\n",
    "for _ in range(n_trials):\n",
    "    simulate_single_trial(n_steps=100)\n",
    "time_sequential = time.time() - start\n",
    "\n",
    "# Batched trials\n",
    "start = time.time()\n",
    "simulate_batched_trials(n_trials=n_trials, n_steps=100)\n",
    "time_batched = time.time() - start\n",
    "\n",
    "print(f\"⏱️  Sequential (10 trials): {time_sequential:.3f} seconds\")\n",
    "print(f\"⏱️  Batched (10 trials): {time_batched:.3f} seconds\")\n",
    "print(f\"\\n🚀 Batching speedup: {time_sequential/time_batched:.1f}×\")\n",
    "print(f\"\\n💡 Batching is especially effective on GPUs!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: GPU Acceleration\n",
    "\n",
    "GPUs excel at parallel operations on large arrays. BrainPy automatically uses GPUs when available via JAX.\n",
    "\n",
    "**GPU benefits:**\n",
    "- Massive parallelism (1000s of cores)\n",
    "- High memory bandwidth\n",
    "- Fast matrix operations\n",
    "- 10-100× speedup for large networks\n",
    "\n",
    "**Best practices:**\n",
    "- Use large batch sizes\n",
    "- Minimize CPU-GPU data transfer\n",
    "- Keep data on GPU between operations\n",
    "- Use JIT compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if GPU is available\n",
    "try:\n",
    "    gpu_device = jax.devices('gpu')[0]\n",
    "    has_gpu = True\n",
    "    print(\"✅ GPU detected:\", gpu_device)\n",
    "except:\n",
    "    has_gpu = False\n",
    "    print(\"ℹ️  No GPU detected, using CPU\")\n",
    "\n",
    "if has_gpu:\n",
    "    # Compare CPU vs GPU for large operation\n",
    "    n = 10000\n",
    "    \n",
    "    # CPU\n",
    "    with jax.default_device(jax.devices('cpu')[0]):\n",
    "        x = jax.random.normal(jax.random.PRNGKey(0), (n, n))\n",
    "        \n",
    "        start = time.time()\n",
    "        y = jnp.dot(x, x)\n",
    "        y.block_until_ready()  # Wait for computation\n",
    "        time_cpu = time.time() - start\n",
    "    \n",
    "    # GPU\n",
    "    with jax.default_device(gpu_device):\n",
    "        x = jax.random.normal(jax.random.PRNGKey(0), (n, n))\n",
    "        \n",
    "        start = time.time()\n",
    "        y = jnp.dot(x, x)\n",
    "        y.block_until_ready()\n",
    "        time_gpu = time.time() - start\n",
    "    \n",
    "    print(f\"\\n🖥️  CPU time: {time_cpu:.4f} seconds\")\n",
    "    print(f\"🎮 GPU time: {time_gpu:.4f} seconds\")\n",
    "    print(f\"🚀 GPU speedup: {time_cpu/time_gpu:.1f}×\")\n",
    "else:\n",
    "    print(\"\\n💡 To use GPU:\")\n",
    "    print(\"   1. Install JAX with GPU support\")\n",
    "    print(\"   2. Install CUDA drivers\")\n",
    "    print(\"   3. BrainPy will automatically use GPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6: Performance Profiling\n",
    "\n",
    "To optimize performance, you need to identify bottlenecks. Use profiling to find where time is spent.\n",
    "\n",
    "**Profiling strategies:**\n",
    "1. **Time individual operations**: Find slow components\n",
    "2. **Use JAX profiler**: Detailed GPU/TPU profiling\n",
    "3. **Monitor memory**: Detect memory leaks\n",
    "4. **Check compilation**: Ensure JIT is working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple profiling example\n",
    "class ProfilingNetwork(brainstate.nn.Module):\n",
    "    def __init__(self, n_neurons=5000):\n",
    "        super().__init__()\n",
    "        self.lif = brainpy.state.LIF(n_neurons, V_rest=-65.*u.mV, V_th=-50.*u.mV, tau=10.*u.ms)\n",
    "        self.proj = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_neurons, n_neurons, 0.01, 0.5*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_neurons, tau=5.*u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.lif\n",
    "        )\n",
    "    \n",
    "    def update(self, inp):\n",
    "        spk = self.lif.get_spike()\n",
    "        self.proj(spk)\n",
    "        self.lif(inp)\n",
    "        return spk\n",
    "\n",
    "# Profile simulation\n",
    "net = ProfilingNetwork(n_neurons=5000)\n",
    "brainstate.nn.init_all_states(net)\n",
    "\n",
    "@brainstate.transform.jit\n",
    "def run_step(net, inp):\n",
    "    return net(inp)\n",
    "\n",
    "# Warmup\n",
    "inp = brainstate.random.rand(5000) * 2.0 * u.nA\n",
    "_ = run_step(net, inp)\n",
    "\n",
    "# Profile multiple steps\n",
    "n_steps = 100\n",
    "step_times = []\n",
    "\n",
    "for _ in range(n_steps):\n",
    "    inp = brainstate.random.rand(5000) * 2.0 * u.nA\n",
    "    \n",
    "    start = time.time()\n",
    "    _ = run_step(net, inp)\n",
    "    step_times.append(time.time() - start)\n",
    "\n",
    "step_times = np.array(step_times) * 1000  # Convert to ms\n",
    "\n",
    "# Plot timing distribution\n",
    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
    "\n",
    "# Time series\n",
    "axes[0].plot(step_times, 'b-', linewidth=1, alpha=0.7)\n",
    "axes[0].axhline(np.mean(step_times), color='r', linestyle='--', \n",
    "                label=f'Mean: {np.mean(step_times):.2f} ms')\n",
    "axes[0].set_xlabel('Step', fontsize=12)\n",
    "axes[0].set_ylabel('Time (ms)', fontsize=12)\n",
    "axes[0].set_title('Step-by-Step Timing', fontsize=14, fontweight='bold')\n",
    "axes[0].legend()\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Histogram\n",
    "axes[1].hist(step_times, bins=30, color='blue', alpha=0.7, edgecolor='black')\n",
    "axes[1].axvline(np.mean(step_times), color='r', linestyle='--', linewidth=2,\n",
    "               label=f'Mean: {np.mean(step_times):.2f} ms')\n",
    "axes[1].set_xlabel('Time (ms)', fontsize=12)\n",
    "axes[1].set_ylabel('Frequency', fontsize=12)\n",
    "axes[1].set_title('Timing Distribution', fontsize=14, fontweight='bold')\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, alpha=0.3, axis='y')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"📊 Performance Statistics:\")\n",
    "print(f\"   Mean time/step: {np.mean(step_times):.2f} ms\")\n",
    "print(f\"   Std deviation: {np.std(step_times):.2f} ms\")\n",
    "print(f\"   Min time: {np.min(step_times):.2f} ms\")\n",
    "print(f\"   Max time: {np.max(step_times):.2f} ms\")\n",
    "print(f\"   Throughput: {1000/np.mean(step_times):.1f} steps/second\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 7: Optimization Checklist\n",
    "\n",
    "Here's a comprehensive checklist for optimizing large-scale simulations.\n",
    "\n",
    "### Before Optimization\n",
    "1. **Profile first**: Identify actual bottlenecks\n",
    "2. **Set target**: Define performance goals\n",
    "3. **Baseline**: Measure current performance\n",
    "\n",
    "### Code Optimizations\n",
    "- ✅ Use JIT compilation (`@brainstate.transform.jit`)\n",
    "- ✅ Use sparse connectivity (`EventFixedProb`)\n",
    "- ✅ Use float32 instead of float64\n",
    "- ✅ Batch multiple trials together\n",
    "- ✅ Avoid Python loops (use `for_loop` or `scan`)\n",
    "- ✅ Minimize state storage\n",
    "- ✅ Use appropriate time steps (larger = faster)\n",
    "\n",
    "### Hardware Optimizations\n",
    "- ✅ Use GPU/TPU when available\n",
    "- ✅ Increase batch size for better GPU utilization\n",
    "- ✅ Monitor GPU memory usage\n",
    "- ✅ Keep data on accelerator (avoid CPU-GPU transfers)\n",
    "\n",
    "### Algorithm Optimizations\n",
    "- ✅ Simplify neuron models if possible\n",
    "- ✅ Use event-driven dynamics where appropriate\n",
    "- ✅ Reduce synaptic computations (sparse updates)\n",
    "- ✅ Cache frequently computed values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstrate optimization impact\n",
    "def benchmark_configurations():\n",
    "    \"\"\"Benchmark different optimization strategies.\"\"\"\n",
    "    \n",
    "    n_neurons = 2000\n",
    "    n_steps = 100\n",
    "    results = {}\n",
    "    \n",
    "    # 1. Baseline (no optimizations)\n",
    "    print(\"Testing: Baseline (no optimizations)...\")\n",
    "    net1 = SimpleNetwork(n_neurons)\n",
    "    brainstate.nn.init_all_states(net1)\n",
    "    \n",
    "    start = time.time()\n",
    "    for _ in range(n_steps):\n",
    "        inp = brainstate.random.rand(n_neurons) * 2.0 * u.nA\n",
    "        _ = net1(inp)\n",
    "    results['Baseline'] = time.time() - start\n",
    "    \n",
    "    # 2. With JIT\n",
    "    print(\"Testing: With JIT...\")\n",
    "    net2 = SimpleNetwork(n_neurons)\n",
    "    brainstate.nn.init_all_states(net2)\n",
    "    \n",
    "    @brainstate.transform.jit\n",
    "    def step_jit(net, inp):\n",
    "        return net(inp)\n",
    "    \n",
    "    # Warmup\n",
    "    inp = brainstate.random.rand(n_neurons) * 2.0 * u.nA\n",
    "    _ = step_jit(net2, inp)\n",
    "    \n",
    "    start = time.time()\n",
    "    for _ in range(n_steps):\n",
    "        inp = brainstate.random.rand(n_neurons) * 2.0 * u.nA\n",
    "        _ = step_jit(net2, inp)\n",
    "    results['JIT'] = time.time() - start\n",
    "    \n",
    "    # 3. With JIT + Batching\n",
    "    print(\"Testing: JIT + Batching...\")\n",
    "    batch_size = 10\n",
    "    net3 = SimpleNetwork(n_neurons)\n",
    "    brainstate.nn.init_all_states(net3, batch_size=batch_size)\n",
    "    \n",
    "    # Warmup\n",
    "    inp = brainstate.random.rand(batch_size, n_neurons) * 2.0 * u.nA\n",
    "    _ = step_jit(net3, inp)\n",
    "    \n",
    "    start = time.time()\n",
    "    for _ in range(n_steps):\n",
    "        inp = brainstate.random.rand(batch_size, n_neurons) * 2.0 * u.nA\n",
    "        _ = step_jit(net3, inp)\n",
    "    results['JIT+Batch'] = time.time() - start\n",
    "    \n",
    "    return results\n",
    "\n",
    "# Run benchmark\n",
    "results = benchmark_configurations()\n",
    "\n",
    "# Visualize results\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "configs = list(results.keys())\n",
    "times = list(results.values())\n",
    "speedups = [times[0] / t for t in times]\n",
    "\n",
    "bars = ax.bar(configs, times, color=['red', 'orange', 'green'], alpha=0.7)\n",
    "\n",
    "# Add speedup labels\n",
    "for i, (bar, speedup) in enumerate(zip(bars, speedups)):\n",
    "    height = bar.get_height()\n",
    "    ax.text(bar.get_x() + bar.get_width()/2., height,\n",
    "           f'{speedup:.1f}× faster\\n{times[i]:.2f}s',\n",
    "           ha='center', va='bottom', fontsize=11, fontweight='bold')\n",
    "\n",
    "ax.set_ylabel('Time (seconds)', fontsize=12)\n",
    "ax.set_title('Optimization Impact (2000 neurons, 100 steps)', fontsize=14, fontweight='bold')\n",
    "ax.grid(True, alpha=0.3, axis='y')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"\\n📊 Optimization Results:\")\n",
    "for config, t in results.items():\n",
    "    speedup = times[0] / t\n",
    "    print(f\"   {config:15s}: {t:.3f}s  ({speedup:.1f}× faster)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 8: Complete Large-Scale Example\n",
    "\n",
    "Let's put it all together with a fully optimized large-scale simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimized large-scale network\n",
    "class OptimizedLargeNetwork(brainstate.nn.Module):\n",
    "    \"\"\"Fully optimized large-scale E-I network.\"\"\"\n",
    "    \n",
    "    def __init__(self, n_exc=8000, n_inh=2000, p_conn=0.02):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_exc = n_exc\n",
    "        self.n_inh = n_inh\n",
    "        \n",
    "        # LIF neurons (using default float32)\n",
    "        self.E = brainpy.state.LIF(n_exc, V_rest=-65.*u.mV, V_th=-50.*u.mV, \n",
    "                       V_reset=-65.*u.mV, tau=15.*u.ms)\n",
    "        self.I = brainpy.state.LIF(n_inh, V_rest=-65.*u.mV, V_th=-50.*u.mV,\n",
    "                       V_reset=-65.*u.mV, tau=10.*u.ms)\n",
    "        \n",
    "        # Sparse connectivity\n",
    "        self.E2E = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_exc, n_exc, p_conn, 0.05*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_exc, tau=5.*u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.E\n",
    "        )\n",
    "        \n",
    "        self.E2I = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_exc, n_inh, p_conn, 0.05*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_inh, tau=5.*u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.I\n",
    "        )\n",
    "        \n",
    "        self.I2E = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_inh, n_exc, p_conn, 0.4*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_exc, tau=10.*u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.E\n",
    "        )\n",
    "        \n",
    "        self.I2I = brainpy.state.AlignPostProj(\n",
    "            comm=brainstate.nn.EventFixedProb(n_inh, n_inh, p_conn, 0.4*u.mS),\n",
    "            syn=brainpy.state.Expon.desc(n_inh, tau=10.*u.ms),\n",
    "            out=brainpy.state.CUBA.desc(),\n",
    "            post=self.I\n",
    "        )\n",
    "    \n",
    "    def update(self, inp_e, inp_i):\n",
    "        # Get spikes\n",
    "        spk_e = self.E.get_spike()\n",
    "        spk_i = self.I.get_spike()\n",
    "        \n",
    "        # Update projections\n",
    "        self.E2E(spk_e)\n",
    "        self.E2I(spk_e)\n",
    "        self.I2E(spk_i)\n",
    "        self.I2I(spk_i)\n",
    "        \n",
    "        # Update neurons\n",
    "        self.E(inp_e)\n",
    "        self.I(inp_i)\n",
    "        \n",
    "        return spk_e, spk_i\n",
    "\n",
    "# Create and simulate\n",
    "print(\"Creating large-scale network...\")\n",
    "large_net = OptimizedLargeNetwork(n_exc=8000, n_inh=2000, p_conn=0.02)\n",
    "brainstate.nn.init_all_states(large_net)\n",
    "\n",
    "print(\"\\n📊 Network Statistics:\")\n",
    "print(f\"   Total neurons: {large_net.n_exc + large_net.n_inh:,}\")\n",
    "print(f\"   Excitatory: {large_net.n_exc:,} (80%)\")\n",
    "print(f\"   Inhibitory: {large_net.n_inh:,} (20%)\")\n",
    "print(f\"   Connectivity: 2%\")\n",
    "print(f\"   Estimated connections: {10000*10000*0.02:,.0f}\")\n",
    "print(f\"   Estimated memory: ~50 MB\")\n",
    "\n",
    "# JIT-compiled simulation\n",
    "@brainstate.transform.jit\n",
    "def simulate_step(net, inp_e, inp_i):\n",
    "    return net(inp_e, inp_i)\n",
    "\n",
    "# Warmup\n",
    "print(\"\\nCompiling (this takes a moment)...\")\n",
    "inp_e = brainstate.random.rand(large_net.n_exc) * 1.0 * u.nA\n",
    "inp_i = brainstate.random.rand(large_net.n_inh) * 1.0 * u.nA\n",
    "_ = simulate_step(large_net, inp_e, inp_i)\n",
    "print(\"✅ Compilation complete!\")\n",
    "\n",
    "# Run simulation\n",
    "print(\"\\nRunning simulation...\")\n",
    "n_steps = 500\n",
    "spike_history_e = []\n",
    "spike_history_i = []\n",
    "\n",
    "start = time.time()\n",
    "for i in range(n_steps):\n",
    "    inp_e = brainstate.random.rand(large_net.n_exc) * 1.0 * u.nA\n",
    "    inp_i = brainstate.random.rand(large_net.n_inh) * 1.0 * u.nA\n",
    "    spk_e, spk_i = simulate_step(large_net, inp_e, inp_i)\n",
    "    \n",
    "    # Downsample recording (save memory)\n",
    "    if i % 5 == 0:\n",
    "        spike_history_e.append(spk_e)\n",
    "        spike_history_i.append(spk_i)\n",
    "\n",
    "sim_time = time.time() - start\n",
    "\n",
    "print(f\"\\n⏱️  Simulation complete:\")\n",
    "print(f\"   Real time: {sim_time:.2f} seconds\")\n",
    "print(f\"   Simulated time: {n_steps * 0.1} ms\")\n",
    "print(f\"   Speedup: {(n_steps * 0.1 / 1000) / sim_time:.1f}× real-time\")\n",
    "print(f\"   Throughput: {n_steps / sim_time:.1f} steps/second\")\n",
    "\n",
    "# Visualize downsampled activity\n",
    "spike_history_e = jnp.array(spike_history_e)\n",
    "spike_history_i = jnp.array(spike_history_i)\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True)\n",
    "\n",
    "# Excitatory raster (subsample neurons for visibility)\n",
    "n_show = 500\n",
    "times_ms = np.arange(len(spike_history_e)) * 5 * 0.1  # Downsampled times\n",
    "\n",
    "for neuron_idx in range(min(n_show, large_net.n_exc)):\n",
    "    spike_times = times_ms[spike_history_e[:, neuron_idx] > 0]\n",
    "    axes[0].scatter(spike_times, [neuron_idx] * len(spike_times),\n",
    "                   s=0.5, c='blue', alpha=0.5)\n",
    "\n",
    "axes[0].set_ylabel('Excitatory Neuron', fontsize=12)\n",
    "axes[0].set_title(f'Large-Scale Network Activity ({large_net.n_exc + large_net.n_inh:,} neurons)', \n",
    "                 fontsize=14, fontweight='bold')\n",
    "axes[0].set_ylim(0, n_show)\n",
    "\n",
    "# Inhibitory raster\n",
    "for neuron_idx in range(large_net.n_inh):\n",
    "    spike_times = times_ms[spike_history_i[:, neuron_idx] > 0]\n",
    "    axes[1].scatter(spike_times, [neuron_idx] * len(spike_times),\n",
    "                   s=0.5, c='red', alpha=0.5)\n",
    "\n",
    "axes[1].set_xlabel('Time (ms)', fontsize=12)\n",
    "axes[1].set_ylabel('Inhibitory Neuron', fontsize=12)\n",
    "axes[1].set_title('Inhibitory Population', fontsize=14, fontweight='bold')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"\\n✅ Successfully simulated 10,000 neuron network!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this tutorial, you learned:\n",
    "\n",
    "✅ **JIT compilation**\n",
    "   - Use `@brainstate.transform.jit` for 10-100× speedup\n",
    "   - Functions must be pure and have static shapes\n",
    "   - Essential for large-scale simulations\n",
    "\n",
    "✅ **Memory optimization**\n",
    "   - Use float32 instead of float64 (2× savings)\n",
    "   - Minimize state storage\n",
    "   - Don't accumulate full histories\n",
    "\n",
    "✅ **Sparse connectivity**\n",
    "   - Use `EventFixedProb` for automatic sparse operations\n",
    "   - 90-99% memory reduction for biological connectivity\n",
    "   - Faster computation (skip zero connections)\n",
    "\n",
    "✅ **Batching**\n",
    "   - Run multiple trials simultaneously\n",
    "   - Better hardware utilization\n",
    "   - Faster parameter sweeps\n",
    "\n",
    "✅ **GPU/TPU acceleration**\n",
    "   - Automatic via JAX when available\n",
    "   - 10-100× speedup for large networks\n",
    "   - Keep data on device\n",
    "\n",
    "✅ **Performance profiling**\n",
    "   - Identify bottlenecks before optimizing\n",
    "   - Monitor memory usage\n",
    "   - Track throughput metrics\n",
    "\n",
    "**Optimization workflow:**\n",
    "\n",
    "```python\n",
    "# 1. Create network with sparse connectivity\n",
    "net = OptimizedNetwork(\n",
    "    n_neurons=10000,\n",
    "    connectivity=0.02  # Sparse!\n",
    ")\n",
    "\n",
    "# 2. Initialize with batching\n",
    "brainstate.nn.init_all_states(net, batch_size=10)\n",
    "\n",
    "# 3. JIT compile simulation loop\n",
    "@brainstate.transform.jit\n",
    "def simulate_step(net, inp):\n",
    "    return net(inp)\n",
    "\n",
    "# 4. Run on GPU (automatic if available)\n",
    "for i in range(n_steps):\n",
    "    inp = get_input()\n",
    "    output = simulate_step(net, inp)\n",
    "```\n",
    "\n",
    "**Scale achieved:**\n",
    "- ✅ 10,000 neurons: Easy on CPU\n",
    "- ✅ 100,000 neurons: Needs GPU\n",
    "- ✅ 1,000,000+ neurons: Multi-GPU or TPU\n",
    "\n",
    "**Next steps:**\n",
    "- Try your own large-scale models\n",
    "- Experiment with different connectivity patterns\n",
    "- Profile and optimize your specific use case\n",
    "- Use specialized tutorials for specific applications\n",
    "- Explore multi-GPU scaling (advanced)\n",
    "\n",
    "**References:**\n",
    "- JAX documentation: https://jax.readthedocs.io/\n",
    "- BrainPy optimization guide: https://brainpy.readthedocs.io/\n",
    "- Neuromorphic computing benchmarks\n",
    "- Large-scale brain simulation papers (Spaun, Blue Brain Project)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "Test your understanding:\n",
    "\n",
    "### Exercise 1: JIT Compilation\n",
    "Take a non-JIT network and apply JIT compilation. Measure the speedup. What happens if you violate JIT rules (e.g., use Python loops)?\n",
    "\n",
    "### Exercise 2: Memory Analysis\n",
    "Estimate memory requirements for a 100,000 neuron network with 1% connectivity. Will it fit in 16GB RAM?\n",
    "\n",
    "### Exercise 3: Sparse vs Dense\n",
    "Implement the same network with dense and sparse connectivity. Compare memory usage and runtime.\n",
    "\n",
    "### Exercise 4: Batching Strategy\n",
    "Run 100 independent trials. Compare: (a) sequential, (b) batched 10×10, (c) batched 100×1. Which is fastest?\n",
    "\n",
    "### Exercise 5: Profiling\n",
    "Profile a large network and identify the slowest operation. Optimize it and measure improvement.\n",
    "\n",
    "**Bonus Challenge:** Scale up to the largest network your hardware can handle. How many neurons can you simulate in real-time?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ecosystem-py",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
